#! /usr/bin/env python3
# -*-coding: utf-8-*-

"""
Created on Apr 20

@author: Moonkie

@attention: 
"""

import mxnet as mx
import numpy as np
import cv2
import os


CURDIR = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.split(os.path.dirname(os.path.abspath(__file__)))[0]
LOGDIR = os.path.join(CURDIR,'log')
ERROR_LOG_FILE = os.path.join(LOGDIR,'error.log')
KVSTORETYPE = 'local'
LST = os.path.join(CURDIR,'train.lst')

if not os.path.exists(LOGDIR):
    os.makedirs(LOGDIR)

class NetGenerator(object):
    @staticmethod
    def convolutionPoolingFactory(data,name,num_filter,cov_kernel,cov_stride,pool_kernel,act_type='leaky',slope=.25,pool_type='avg',pool_stride=(1,1),cov_pad=(0,0),pool_pad=(0,0),pool_global=False,cov_cudnn=False,lower_bound=.125,upper_bound=.334):
        cov_name = 'conv_%d' % int(name)
        act_name = 'act_%d' % int(name)
        pool_name = 'pool_%d' % int(name)
        
        cov = mx.sym.Convolution(data=data,name=cov_name,kernel=cov_kernel,stride=cov_stride,pad=cov_pad,num_filter=num_filter,cudnn_off=cov_cudnn)

        act = mx.sym.LeakyReLU(cov,act_type=act_type,name=act_name,slope=slope)

        pool = mx.sym.Pooling(data=act,name=pool_name,kernel=pool_kernel,stride=pool_stride,pad=pool_pad,global_pool=pool_global,pool_type=pool_type)
        return pool

    @staticmethod
    def convolutionFactory(data,name,num_filter,cov_kernel,cov_stride,act_type='relu',cov_pad=(0,0),cov_cudnn=False):
        cov_name = 'conv_%d' % int(name) 
        act_name = 'act_%d' % int(name) 

        cov = mx.sym.Convolution(data=data,name=cov_name,num_filter=num_filter,kernel=cov_kernel,stride=cov_kernel,pad=cov_pad,cudnn_off=cov_cudnn)

        act = mx.sym.Activation(data=cov,name=act_name,act_type=act_type)
        return act

    @staticmethod
    def leakyConvFactory(data,name,num_filter,cov_kernel,cov_stride,act_type='leaky',slope=.25,lower_bound=.125,upper_bound=.334,cov_pad=(0,0),cov_cudnn=False):
        cov_name = 'conv_%d' % int(name) 
        act_name = 'act_%d' % int(name) 

        cov = mx.sym.Convolution(data=data,name=cov_name,num_filter=num_filter,kernel=cov_kernel,stride=cov_kernel,pad=cov_pad,cudnn_off=cov_cudnn)

        act = mx.sym.LeakyReLU(data=cov,name=act_name,act_type=act_type,slope=slope,lower_bound=lower_bound,upper_bound=upper_bound)
        return act

    @staticmethod
    def fullyConnectedActivationFactory(data,name,num_hidden,act_type='relu',no_bias=False):
        fc_name = 'fc_%d' % int(name)
        act_name = 'act_%d' % int(name)

        fc = mx.sym.FullyConnected(data=data,name=fc_name,num_hidden=num_hidden,no_bias=no_bias)
        
        act = mx.sym.Activation(data=fc,name=act_name,act_type=act_type)
        return act

    @staticmethod
    def leakyFcFactory(data,name,num_hidden,act_type='leaky',slope=.25,lower_bound=.125,upper_bound=.334,no_bias=False):
        fc_name = 'fc_%d' % int(name)
        act_name = 'act_%d' % int(name)

        fc = mx.sym.FullyConnected(data=data,name=fc_name,num_hidden=num_hidden,no_bias=no_bias)

        act = mx.sym.LeakyReLU(data=fc,name=act_name,act_type=act_type,slope=slope,lower_bound=lower_bound,upper_bound=upper_bound)
        return act

    @staticmethod
    def fcActivationFcFactory(data,name,num_hidden1,num_hidden2,act_type='relu',no_bias=False):
        fc1_name = 'fc_%d_1' % int(name)
        fc2_name = 'fc_%d_2' % int(name)
        act_name = 'act_%d' % int(name)

        fc1 = mx.sym.FullyConnected(data=data,name=fc1_name,num_hidden=num_hidden1,no_bias=no_bias)

        act = mx.sym.Activation(data=fc1,name=act_name,act_type=act_type)

        fc2 = mx.sym.FullyConnected(data=act,name=fc2_name,num_hidden=num_hidden2)
        return fc2

    @staticmethod
    def leakyFcAndFc(data,name,num_hidden1,num_hidden2,act_type='leaky',slope=.25,lower_bound=.125,upper_bound=.334,no_bias=False,firsted=False,data_shape=None):
        fc1_name = 'fc_%d_1' % int(name)
        fc2_name = 'fc_%d_2' % int(name)
        act_name = 'act_%d' % int(name)

        if firsted:
            bias = mx.sym.ones(shape=data_shape,name=fc1_name+'_bias')
            fc1 = mx.sym.FullyConnected(data=data,name=fc1_name,num_hidden=num_hidden1,bias=bias)
        else:
            fc1 = mx.sym.FullyConnected(data=data,name=fc1_name,num_hidden=num_hidden1,no_bias=no_bias)

        act = mx.sym.LeakyReLU(data=fc1,name=act_name,act_type=act_type,slope=slope,lower_bound=lower_bound,upper_bound=upper_bound)

        fc2 = mx.sym.FullyConnected(data=act,name=fc2_name,num_hidden=num_hidden2)
        return fc2


def get_symbol(num_classes):
    data = mx.sym.Variable('data')
    norm = mx.sym.BatchNorm(data,name='norm')

    pool1 = NetGenerator.convolutionPoolingFactory(norm,name=1,num_filter=48,cov_kernel=(3,3),cov_stride=(1,1),pool_kernel=(2,2),pool_stride=(1,1),act_type='prelu',pool_type='max',cov_pad=(1,1),cov_cudnn=True)

    pool2 = NetGenerator.convolutionPoolingFactory(pool1,name=2,num_filter=64,cov_kernel=(2,2),cov_stride=(1,1),pool_kernel=(2,2),pool_stride=(2,2),act_type='prelu',pool_type='max',cov_cudnn=True)

    pool3 = NetGenerator.convolutionPoolingFactory(pool2,name=3,num_filter=80,cov_kernel=(3,3),cov_stride=(1,1),cov_pad=(1,1),pool_kernel=(2,2),pool_stride=(1,1),act_type='prelu',pool_type='max',cov_cudnn=True)

    pool4 = NetGenerator.convolutionPoolingFactory(pool3,name=4,num_filter=96,cov_kernel=(3,3),cov_stride=(1,1),pool_kernel=(2,2),pool_stride=(2,2),act_type='prelu',pool_type='max',cov_cudnn=True)

    pool5 = NetGenerator.convolutionPoolingFactory(pool4,name=5,num_filter=112,cov_kernel=(2,2),cov_stride=(1,1),pool_kernel=(2,2),pool_stride=(1,1),act_type='prelu',pool_type='max',cov_cudnn=True)

    pool6 = NetGenerator.convolutionPoolingFactory(pool5,name=6,num_filter=128,cov_kernel=(2,2),cov_stride=(1,1),pool_kernel=(3,3),pool_stride=(2,2),act_type='prelu',pool_type='max',cov_cudnn=True)

    # avgpool = mx.sym.Pooling(pool4,global_pool=True,name='avg_pool',kernel=(2,2),stride=(1,1),pool_type='avg')
    # fc1 = NetGenerator.fullyConnectedActivationFactory(pool4,name=3,num_hidden=256)
    fc1 = NetGenerator.leakyFcAndFc(pool6,name=7,num_hidden1=512,num_hidden2=256,act_type='prelu')
    drop1 = mx.sym.Dropout(fc1,p=.25,name='drop_1')
    fc2 = mx.sym.FullyConnected(fc1,num_hidden=num_classes,name='fc_2')
    out = mx.sym.SoftmaxOutput(fc2,name='softmax')
    return out

def get_iterator(data_shape,batch_size,
                    trainlst,trainrec,
                    validlst,validrec,
                    trainmean,validmean):

    train = mx.io.ImageRecordIter(
        path_imglist = trainlst,
        path_imgrec = trainrec,
        mean_img    = trainmean,
        data_shape  = data_shape,
        batch_size  = batch_size,
        rand_crop   = False,
        rand_mirror = True,
        round_batch = False)

    val = mx.io.ImageRecordIter(
        path_imglist = validlst,
        path_imgrec = validrec,
        mean_img    = validmean,
        data_shape  = data_shape,
        batch_size  = batch_size,
        rand_crop   = False,
        rand_mirror = False)
    return train,val


def get_test_iterator(data_shape,batch_size,
                    testlst,testrec,testmean):
    '''Test Image data sets generator
    '''
    test = mx.io.ImageRecordIter(
        path_imglist = testlst,
        path_imgrec = testrec,
        mean_img    = testmean,
        data_shape  = data_shape,
        batch_size  = batch_size,
        rand_crop   = False,
        rand_mirror = False)
    return test


def showimg(img):
    cv2.imshow('image',img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == '__main__':
    traindata,valdata,testdata,shape = get_fer_data('trains/fer2013.csv',num_batch=100,channel=1)
    # print(dir(traindata))
    # print(traindata.getdata())
    # data = get_train_iterator((1,64,64),10,'trains/image.rec','trains/image.lst',mean_img='trains/image.bin')
    # d = data.getdata()
    # print(d)
    pass
